import requests
from requests_toolbelt import MultipartEncoder
import os
import contextlib
import mimetypes
import json
import base64
from dotenv import load_dotenv

load_dotenv()

API_URL = "https://ai.gitee.com/v1/images/edits"
api_key = os.getenv("api_key")
headers = {
	"Authorization": f"Bearer {api_key}" # 使用 Bearer 令牌进行身份验证 
}

def query(payload):
	fields = [
		("prompt", payload["prompt"]),
		("model", payload["model"]),
		("size", payload["size"]),
		("seed", str(payload["seed"])),
		("num_inference_steps", str(payload["num_inference_steps"])),
		("guidance_scale", str(payload["guidance_scale"])),
	]
	with contextlib.ExitStack() as stack:
		for item in payload["task_types"]:
			fields.append(("task_types", item if isinstance(item, str) else json.dumps(item)))
		for filepath in payload["image"]:
			name = os.path.basename(filepath)
			if filepath.startswith(("http://", "https://")):
				response = requests.get(filepath, timeout=10)
				response.raise_for_status()
				fields.append(("image", (name, response.content, response.headers.get("Content-Type", "application/octet-stream"))))
			else:
				mime_type, _ = mimetypes.guess_type(filepath)
				fields.append(("image", (name, stack.enter_context(open(filepath, "rb")), mime_type or "application/octet-stream")))
		encoder = MultipartEncoder(fields)
		headers["Content-Type"] = encoder.content_type
		response = requests.post(API_URL, headers=headers, data=encoder)
		return response.json()

output = query({
	"prompt": "profile shot dark photo of a 25-year-old female with smoke escaping from her mouth, the backlit smoke gives the image an ephemeral quality, natural face, natural eyebrows, natural skin texture, award winning photo, highly detailed face, atmospheric lighting, film grain, monochrome",
	"image": [ 
		# 替换为实际的图片路径
		"/path/to/image-id.png",
		"/path/to/image-style.png"
	],
	"task_types": [ 
		# 替换为实际的任务类型：ip（从输入图片中删除背景）、id（从输入图片中提取面部区域）、style（背景保留，prompt 需要以 "generate a same image" 开头）
		"ip",
		"style"
	],
	"model": "DreamO",
	"size": "1024x1024",
	"seed": -1,
	"num_inference_steps": 12,
	"guidance_scale": 3.5
})

for i, image_data in enumerate(output["data"]):
	if "url" in image_data:
		# 从网址下载
		ext = image_data["url"].split('.')[-1].split('?')[0] or "jpg"
		filename = f"DreamO-output-{i}.{ext}"
		response = requests.get(image_data["url"], timeout=30)
		response.raise_for_status()
		with open(filename, "wb") as f:
			f.write(response.content)
		print(f"Downloaded image to {filename}")
	elif "b64_json" in image_data:
		# 解码 Base64
		image_bytes = base64.b64decode(image_data["b64_json"])
		filename = f"DreamO-output-{i}.jpg"
		with open(filename, "wb") as f:
			f.write(image_bytes)
		print(f"Saved image to {filename}")
